from PIL import Image, ImageDraw, ImageFont
import os
import argparse
from datetime import datetime


# 新增函数：解析命令行参数
def parse_arguments():
    """解析命令行参数，允许用户设置总题目数量和 OMR 类型。"""
    parser = argparse.ArgumentParser(
        description="生成可配置的 OMR 答题卡 (矩形选项框) (支持命令行参数配置题目数量和选项类型)。"
    )

    # --questions 参数：设置题目数量，默认值为 10
    parser.add_argument(
        "--questions", type=int, default=10, help="设置答题卡的总题目数量 (默认: 10 题)。"
    )

    # --type 参数：设置 OMR 选项类型，默认值为 SINGLE_CHOICE
    parser.add_argument(
        "--type",
        type=str,
        default="SINGLE_CHOICE",
        choices=["SINGLE_CHOICE", "MULTIPLE_CHOICE", "TRUE_FALSE"],
        help="设置 OMR 选项类型 (SINGLE_CHOICE, MULTIPLE_CHOICE, TRUE_FALSE)。",
    )
    return parser.parse_args()


# 新增函数：加载特定字体
def load_specific_font(font_paths, size, fallback_name):
    """遍历路径列表，尝试加载特定字体，失败则返回默认字体。"""
    selected_path = None

    # 1. 查找存在的字体路径
    for path in font_paths:
        if os.path.exists(path):
            selected_path = path
            break

    # 2. 尝试加载字体
    if selected_path:
        try:
            return ImageFont.truetype(selected_path, size)
        except IOError as e:
            # 字体文件存在但加载失败（可能是损坏或格式问题）
            print(
                f"Warning: 尝试加载 {selected_path} 字体时发生 IO 错误。正在使用默认字体。错误: {e}"
            )
            return ImageFont.load_default()
    else:
        # 字体文件路径不存在
        print(
            f"Warning: 未找到 {fallback_name} 字体。正在使用默认字体，可能影响显示效果。"
        )
        return ImageFont.load_default()


def main():
    # --- 解析命令行参数 ---
    args = parse_arguments()

    # --- 用户配置区 (易于修改) ---
    total_questions = args.questions  # 从命令行参数获取题目数量，默认值为 10
    omr_type = args.type  # 从命令行参数获取 OMR 类型

    # 预定义选项集
    OPTION_SETS = {
        "SINGLE_CHOICE": ["A", "B", "C", "D"],  # 适用于单选题
        "MULTIPLE_CHOICE": ["A", "B", "C", "D", "E"],  # 适用于多选题 (假设最多5个选项)
        "TRUE_FALSE": ["T", "F"],  # 适用于判断题
    }

    # 根据配置获取选项
    options_per_question = OPTION_SETS.get(omr_type, OPTION_SETS["SINGLE_CHOICE"])

    # 确保选项数量至少为 1
    if not options_per_question:
        options_per_question = ["A"]

    # === 布局参数 ===
    margin = 50
    max_cols = 4
    row_height = 70

    # 选项矩形尺寸
    rect_width = 40
    rect_height = 30

    # 选项中心间距
    option_spacing = 50
    # 单题总宽度，会根据选项数量动态调整
    col_spacing = len(options_per_question) * option_spacing + 80

    # === 关键定位块尺寸参数 (与四角定位块尺寸同步) ===
    mark_size = 20  # 四角定位块的尺寸
    marker_width = mark_size  # 姓名/学号定位块的宽度
    marker_height = mark_size  # 姓名/学号定位块的高度
    marker_gap_x = 10  # 定位块与文本之间的水平间隔

    # === 计算图像尺寸 ===
    num_rows = (total_questions + max_cols - 1) // max_cols
    # 图像高度：上下边距 + 姓名/学号行高度 + 题目行数 * 行高
    img_width = margin * 2 + max_cols * col_spacing
    img_height = margin + 60 + num_rows * row_height + margin

    # === 创建图像 ===
    img = Image.new("RGB", (img_width, img_height), "white")
    draw = ImageDraw.Draw(img)

    # === 字体加载 (双字体配置：中文-宋体, 英文/数字-Times New Roman) ===
    font_size = 22

    # --- 1. 宋体/中文字体路径 (label_font) ---
    chinese_font_paths = [
        # macOS: Songti SC Regular
        "/System/Library/Fonts/Supplemental/Songti.ttc",
        # Windows: SimSun (宋体), SimHei (黑体 - 作为宋体失败的中文回退)
        "C:/Windows/Fonts/simsun.ttc",
        "C:/Windows/Fonts/simhei.ttf",
        # Linux: Noto CJK Serif (作为宋体替代)
        "/usr/share/fonts/truetype/noto/NotoSerifCJK-Regular.ttc",
    ]
    label_font = load_specific_font(chinese_font_paths, font_size, "宋体/中文字体")

    # --- 2. Times New Roman/英文字母和数字字体路径 (option_font) ---
    latin_font_paths = [
        # Windows: Times New Roman
        "C:/Windows/Fonts/times.ttf",
        "C:/Windows/Fonts/timesbd.ttf",
        # macOS: Times New Roman
        "/System/Library/Fonts/Supplemental/Times New Roman.ttf",
        "/System/Library/Fonts/Times.ttc",
        # Linux: Times New Roman (或通用衬线字体)
        "/usr/share/fonts/truetype/msttcorefonts/Times_New_Roman.ttf",
        "/usr/share/fonts/truetype/dejavu/DejaVuSerif.ttf",
    ]
    option_font = load_specific_font(
        latin_font_paths, font_size, "Times New Roman/英文字体"
    )

    # === 1. 姓名 + 学号（独立定位块且几何中心对齐）===
    text_y_start = margin

    name_text_full = "姓名：____________________"
    student_id_text_full = "学号：____________________"

    # --- 计算中文文本的几何中心Y坐标 (Text Center Y) ---
    try:
        # 使用 label_font (宋体) 获取精确文本边界用于垂直居中对齐
        name_label_bbox = label_font.getbbox("姓名：")
    except AttributeError:
        # 旧版 Pillow 兼容
        name_label_bbox = draw.textbbox((0, 0), "姓名：", font=label_font)

    # 计算文本区域的几何中心Y坐标
    text_center_y = text_y_start + (name_label_bbox[1] + name_label_bbox[3]) / 2

    # --- 姓名区域定位块 ---

    # 定位块 X 坐标：位于文本左侧 (margin - 间隙 - 宽度)
    name_marker_x = margin - marker_width - marker_gap_x

    # 定位块顶部 Y 坐标：使其几何中心与 text_center_y 水平对齐
    name_marker_y_top = text_center_y - marker_height / 2

    draw.rectangle(
        (
            name_marker_x,
            name_marker_y_top,
            name_marker_x + marker_width,
            name_marker_y_top + marker_height,
        ),
        fill="black",
    )

    # 绘制姓名部分 (使用 label_font)
    draw.text((margin, text_y_start), name_text_full, fill="black", font=label_font)

    # 计算学号的起始 X 坐标
    name_area_width_estimate = 400
    gap_between_fields = 60
    student_id_x = margin + name_area_width_estimate + gap_between_fields

    # --- 学号区域定位块 ---

    # 定位块 X 坐标：位于学号文本左侧
    student_id_marker_x = student_id_x - marker_width - marker_gap_x

    # 绘制定位块 (Y 坐标与姓名定位块一致，实现水平对齐)
    draw.rectangle(
        (
            student_id_marker_x,
            name_marker_y_top,
            student_id_marker_x + marker_width,
            name_marker_y_top + marker_height,
        ),
        fill="black",
    )

    # 绘制学号部分 (使用 label_font)
    draw.text(
        (student_id_x, text_y_start),
        student_id_text_full,
        fill="black",
        font=label_font,
    )

    # === 2. 四角定位块（用于页面对齐）===
    # 左上角
    draw.rectangle((0, 0, mark_size, mark_size), fill="black")
    # 右上角
    draw.rectangle((img_width - mark_size, 0, img_width, mark_size), fill="black")
    # 左下角
    draw.rectangle((0, img_height - mark_size, mark_size, img_height), fill="black")
    # 右下角
    draw.rectangle(
        (img_width - mark_size, img_height - mark_size, img_width, img_height),
        fill="black",
    )

    # === 3. 绘制题目（矩形选项框）===
    start_y = margin + 60 + 20  # 姓名/学号字段下方起始 Y 坐标

    for q_num in range(1, total_questions + 1):
        row = (q_num - 1) // max_cols
        col = (q_num - 1) % max_cols

        base_x = margin + col * col_spacing
        base_y = start_y + row * row_height

        # 题号 (使用 option_font - Times New Roman)
        draw.text((base_x, base_y - 10), f"{q_num}.", fill="black", font=option_font)

        # 选项矩形 + 居中文字
        for j, opt in enumerate(options_per_question):
            cx = base_x + 60 + j * option_spacing
            cy = base_y

            # 绘制矩形
            left = cx - rect_width // 2
            top = cy - rect_height // 2
            right = cx + rect_width // 2
            bottom = cy + rect_height // 2
            draw.rectangle((left, top, right, bottom), outline="black", width=2)

            # 选项文字居中 (使用 option_font - Times New Roman)
            try:
                bbox = option_font.getbbox(opt)
            except AttributeError:
                bbox = draw.textbbox((0, 0), opt, font=option_font)

            text_w = bbox[2] - bbox[0]
            text_x = cx - text_w // 2
            # 垂直居中
            text_y = cy - (bbox[1] + bbox[3]) / 2

            draw.text((text_x, text_y), opt, fill="black", font=option_font)

    # === 保存 ===
    # 使用日期时间戳、题目数量和类型生成独特的文件名，并加入形状信息 'rect'
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    output_path = f"omr_card_Q{total_questions}_{omr_type}_rect_{timestamp}.png"

    img.save(output_path)
    print(f"矩形答题卡已成功生成：{os.path.abspath(output_path)}")


if __name__ == "__main__":
    main()
